from scapy.all import IP, UDP, ICMP, TCP, sr1
import asyncio
from functools import partial
import time


# 异步UDP探测函数
async def async_udp_probe(target_ip, ttl):
    packet = IP(dst=target_ip, ttl=ttl) / UDP(dport=33434 + ttl)
    # 使用partial创建一个带有固定参数的函数
    sr1_with_params = partial(sr1, packet, timeout=1, verbose=0)

    try:
        # 将sr1_with_params函数的调用包装在一个线程中运行，并设置超时
        response_future = asyncio.get_event_loop().run_in_executor(None, sr1_with_params)
        response = await asyncio.wait_for(response_future, timeout=1)  # 设置1秒超时
    except (asyncio.TimeoutError, Exception):
        return ttl, '*'  # 超时或错误都表示无响应

    if response is None:
        return ttl, '*'
    else:
        return ttl, response.src


# 异步TCP探测函数
async def async_tcp_probe(target_ip, ttl):
    packet = IP(dst=target_ip, ttl=ttl) / TCP(dport=80 + ttl, flags='S')  # 以常见的80端口且设置SYN标志位为例，可根据实际调整
    # 使用partial创建一个带有固定参数的函数
    sr1_with_params = partial(sr1, packet, timeout=1, verbose=0)

    try:
        # 将sr1_with_params函数的调用包装在一个线程中运行，并设置超时
        response_future = asyncio.get_event_loop().run_in_executor(None, sr1_with_params)
        response = await asyncio.wait_for(response_future, timeout=1)  # 设置1秒超时
    except (asyncio.TimeoutError, Exception):
        return ttl, '*'

    if response is None:
        return ttl, '*'
    else:
        return ttl, response.src


# 异步ICMP探测函数
async def async_icmp_probe(target_ip, ttl, icmp_id):
    packet = IP(dst=target_ip, ttl=ttl) / ICMP(id=icmp_id, seq=ttl)

    # 使用partial创建一个带有固定参数的函数
    sr1_with_params = partial(sr1, packet, timeout=1, verbose=0)

    try:
        # 将sr1_with_params函数的调用包装在一个线程中运行，并设置超时
        response_future = asyncio.get_event_loop().run_in_executor(None, sr1_with_params)
        response = await asyncio.wait_for(response_future, timeout=2)  # 设置1秒超时
    except (asyncio.TimeoutError, Exception):
        return ttl, '*'  # 超时或错误都表示无响应

    if response is None:
        return ttl, '*'
    else:
        # 检查响应类型
        if response.haslayer(ICMP):
            icmp_layer = response[ICMP]
            if icmp_layer.type == 0 and icmp_layer.id == icmp_id:  # ICMP Echo Reply
                return ttl, response.src  # 到达目标
            elif icmp_layer.type == 11 and icmp_layer.code == 0:  # ICMP Time Exceeded
                return ttl, response.src  # 中间路由器
        return ttl, '*'  # 其他类型的ICMP消息或无法识别的响应


# 根据协议选择进行路径追踪的函数，并整合路径，对单个目标的多种协议探测采用异步
async def trace_path(target_ip, max_hops):
    udp_tasks = [async_udp_probe(target_ip, ttl) for ttl in range(1, max_hops + 1)]
    icmp_tasks = [async_icmp_probe(target_ip, ttl, icmp_id=ttl) for ttl in range(1, max_hops + 1)]
    tcp_tasks = [async_tcp_probe(target_ip, ttl) for ttl in range(1, max_hops + 1)]

    # 并发执行UDP、ICMP和TCP任务
    results = await asyncio.gather(*udp_tasks, *icmp_tasks, *tcp_tasks, return_exceptions=True)

    # 初始化整合后的路径列表
    hop_new = ['*'] * max_hops

    # 分别处理UDP、ICMP和TCP的结果，整合路径 对结果results切片操作
    udp_results = results[:max_hops]
    icmp_results = results[max_hops:2 * max_hops]
    tcp_results = results[2 * max_hops:]
    for i in range(len(udp_results)):
        udp_ttl, udp_ip = udp_results[i]
        icmp_ttl, icmp_ip = icmp_results[i]
        tcp_ttl, tcp_ip = tcp_results[i]

        if udp_ip != '*' or icmp_ip != '*':
            if udp_ip != '*':
                hop_new[i] = udp_ip
            else:
                hop_new[i] = icmp_ip
        else:
            if tcp_ip != '*':
                hop_new[i] = tcp_ip
            else:
                hop_new[i] = '*'

    return hop_new


# async def main():
#     targets_list = [
#         ('192.58.128.30', 14),
#         # ('211.95.32.146', 11),  # 目标IP地址和最大跳数
#         # ('111.33.167.71', 12),
#         # ('59.82.44.240', 14),
#         # ('183.47.108.118', 17),
#         # ('106.63.15.10', 16)
#     ]

#     # 对每个目标进行路径探测并整合路径
#     for target_ip, max_hops in targets_list:
#         hop_new = await trace_path(target_ip, max_hops)
#         print(f"目标 {target_ip} 整合后的路径列表: {hop_new}")


# # 运行主协程
# if __name__ == "__main__":
#     start_time = time.time()
#     asyncio.run(main())
#     end_time = time.time()
#     print(f"程序运行耗时: {end_time - start_time} 秒")
